#Import the Module of MetaGPT
import re
import fire
from metagpt.actions import Action, UserRequirement
from metagpt.logs import logger
from metagpt.roles import Role
from metagpt.schema import Message
from metagpt.team import Team
import asyncio
import requests
#Import the Module of Text2Image
from text2image import textTOimage
#Import the Module of Qwen
from Qwen import image_to_text
#Import the Module of download
from download import download_image
#正则表达式，提取prompt部分
def parse_prompt(rsp):
    #prompt 部分需要满足格式 !!prompt prompt_text !!
    pattern = r"!!prompt(.*)!!"
    match = re.search(pattern, rsp , re.DOTALL)
    prompt_text = match.group(1).strip() if match else None
    #return prompt_text -> str
    return prompt_text

#WritePrompts -> Action 用来书写prompt
class WritePrompts(Action):
    #当第一次运行 或 无参考时 in--(instruction:str)
    PROMPT_TEMPLATE:str = """
    你是一个专业的AI提示词工程师，你现在需要书写一个图片生成提示词，
    提示词要用英文词组的形式给出，用逗号隔开，
    你需要遵循以下格式撰写提示词：
    清晰度与质量：如 masterpiece, best quality, 8k, ultra-detailed
    风格：如 realistic, anime, digital painting, cinematic lighting
    主体与数量：如 1girl, 1boy, group of people, animal, robot
    服饰与特征：如 blue dress, school uniform, ponytail, freckles, glasses
    动作与姿态：如 standing, sitting, looking at viewer, hand on hip
    背景与场景：如 indoors, sunset, cyberpunk city, forest, bedroom
    其它部分。
    生成的提示词仅为英文词组的集合，你无需给出任何中文描述，只需要返回提示词，不要有任何其他解释，无需换行
    图片的要求是： {instruction}
    严格围绕instruction，不要偏离主题
    Return !!prompt 这里写你的图片提示词 !! ，不要有任何其它的解释
    你的提示词:
    """
    #第二轮 或 有参考时 in--(instruction:str,temp_info:str,suggestions:str)
    PROMPT_TEMPLATE_OPT:str = """
    你是一个专业的AI提示词工程师，你现在需要书写一个图片生成的提示词，
    提示词要用英文词组的形式给出，用逗号隔开，
    你需要遵循以下格式撰写提示词：
    清晰度与质量：如 masterpiece, best quality, 8k, ultra-detailed
    风格：如 realistic, anime, digital painting, cinematic lighting
    主体与数量：如 1girl, 1boy, group of people, animal, robot
    服饰与特征：如 blue dress, school uniform, ponytail, freckles, glasses
    动作与姿态：如 standing, sitting, looking at viewer, hand on hip
    背景与场景：如 indoors, sunset, cyberpunk city, forest, bedroom
    其它部分。
    图片的要求是 {instruction}
    严格围绕instruction，不要偏离主题
    Return !!prompt 这里写你的图片提示词 !! ，不要有任何其它的解释
    你已经写过一次提示词了，现在需要对提示词进行优化，优化后的提示词需要比之前的提示词更好，之前的提示词是 {temp_info}
    这是给你的修改建议 {suggestions}
    你的提示词:
    """

    #set some basic information
    name:str = "WritePrompts"
    async def run(self,instruction:str,temp_info:str,suggestions:str,is_ref:bool=False)->str:
        is_opt = bool(temp_info)

        if not is_opt:
        #use instruction to generate prompt
            prompt = self.PROMPT_TEMPLATE.format(instruction=instruction)

        else:
            prompt = self.PROMPT_TEMPLATE_OPT.format(instruction=instruction,temp_info=temp_info,suggestions=suggestions)
        print("--------------WritePrompts的aask前的prompt------------------")
        print(prompt)
        print("--------------WritePrompts的aask前的prompt——End------------------")
        rsp = await self._aask(prompt)
        prompt_text = parse_prompt(rsp)
        return prompt_text
    
#PromptWriter -> Role 用来书写prompt
class PromptWriter(Role):
    #set some basic information
        #set name
        name:str = "PromptWriter"
        #set profile of the role 
        profile:str = "PromptWriter"
        def __init__(self,**kwargs):
            super().__init__(**kwargs)

            self._watch([UserRequirement,ProvideSuggestions])

            self.set_actions([WritePrompts])

            #Override the act method
        async def _act(self) -> Message:
            mems = self.get_memories(k=4)

            #information from WritePrompts
            temp_info = next((m.content for m in mems if m.cause_by == '__main__.WritePrompts'), "")
            #information from ProvideSuggestions
            suggestions = next((m.content for m in mems if m.cause_by == '__main__.ProvideSuggestions'), "")

            mems = self.get_memories()
            #information from User
            instruction = next((m.content for m in mems if m.role=="Human"), None)
            #use LLM to generate prompt
            prompt = await self.rc.todo.run(instruction,temp_info,suggestions)
            #add prompt to memory
            msg = Message(content=prompt, role=self.profile, cause_by=type(self.rc.todo))
            self.rc.memory.add(msg)
            #重新添加instruction，防止被刷下去导致下一次查询不到
            #instruction = Message(content=instruction, role="Human", cause_by="UserRequirement")
            #添加instruction到memory
            #self.rc.memory.add(instruction)
            return msg

#GeneratePicture -> Action 用来生成图片
class GeneratePicture(Action):
    #set some basic information
    name:str = "GeneratePicture"

    #Override the run function
    #这里需要调用Text2img的generate_picture函数,Use API from liblib
    async def run(self,prompt:str) -> str:
        picture_url = await asyncio.to_thread(
            textTOimage,
            prompt,
            negative_prompt="ng_deepnegative_v1_75t,(badhandv4:1.2),EasyNegative,(worst quality:2),"
        )
        if not picture_url or "error" in picture_url:
            logger.warning("图片URL可能无效，生成失败？")
            print("图片URL可能无效，生成失败？")
            return None
        #防止图片过大或网速过慢阻塞时间循环，丢入线程池
        await asyncio.to_thread(download_image, picture_url)
        return picture_url

        
#PictureGenerator -> Role 用来生成图片
class PictureGenerator(Role):
    #set some basic information
    name:str = "PictureGenerator"
    profile:str = "PictureGenerator"
    #constructor
    def __init__(self,**kwargs):
        super().__init__(**kwargs)

        self._watch([WritePrompts])
        self.set_actions([GeneratePicture])

    async def _act(self) -> Message:
        #get prompt from memory
        mems = self.get_memories(k=2)
        prompt = next(
            (m.content for m in mems if m.cause_by == '__main__.WritePrompts'), ""
        )

        picture_url = await self.rc.todo.run(prompt)
        #将图片url添加到memory
        msg = Message(content=picture_url, role=self.profile, cause_by=type(self.rc.todo))
        self.rc.memory.add(msg)
        return msg

#ProvideSuggestions -> Action 用来提供建议
class ProvideSuggestions(Action):
    PROMPT_TEMPLATE:str = """
    你是一个专业的AI提示词工程师,你要检查生成的提示词是否完美，并且给出批判性的建议
    这是一张AI生成的图片，用户的要求为：{instruction}
    当前的提示词为：{prompt}
    你要为这个提示词给出修改建议，你无需修改，只需提出修改意见
    Return !!suggestions 你的修改建议 !!
    以 !!suggestions 开始，以 !! 结尾，请勿回复无关回答
    你的修改建议:

    """
    name:str = "ProvideSuggestions"
    async def run(self,instruction:str,prompt:str,picture_url:str) -> str:        
        prompt = self.PROMPT_TEMPLATE.format(instruction=instruction,prompt=prompt)
        suggestions_text = await asyncio.to_thread(
            image_to_text,
            picture_url,
            text_prompt=prompt
        )
        return suggestions_text
    
#SuggestionProvider -> Role 用来提供建议
class SuggestionProvider(Role):

    #set name
    name:str = "SuggestionProvider"
    #set profile of the role
    profile:str = "SuggestionProvider"
    
    def __init__(self,**kwargs):
        super().__init__(**kwargs)

        #set action and watch

        self._watch([GeneratePicture])
        self.set_actions([ProvideSuggestions])

    async def _act(self) -> Message:
    # 从 memory 中获取 prompt（提示词）
        mems = self.get_memories(k=3)
        prompt = next(
            (m.content for m in mems if m.cause_by == '__main__.WritePrompts'), ""
        )
        picture_url = next(
            (m.content for m in mems if m.cause_by == '__main__.GeneratePicture'), ""
        )

        mems = self.get_memories()
        instruction = next(
            (m.content for m in mems if m.role=="Human"), ""
        )


        suggestions = await self.rc.todo.run(instruction,prompt,picture_url)

        print("--------------------------------")
        print(suggestions)
        print("--------------------------------")
        msg = Message(content=suggestions, role=self.profile, cause_by=type(self.rc.todo))
        self.rc.memory.add(msg)

        return msg    

async def main(
    idea :str = "一个女孩，看镜头",
    investment:float = 3.0,
    n_round:int = 10,
    add_human:bool = False
    ):
    
    temp_ref:str = input("Please type y if you need provide a reference,if not , type n (y/n) ")
    TEXT_PROMPT_ALL = '''
        你是一个专业的AI绘画提示词工程师，用户的要求是：{instruction}
        这是用户的参考图，用英文提示词描述这个图片 提示词要用英文词组的形式给出，用逗号隔开，
        你需要遵循以下格式撰写提示词：
        清晰度与质量：如 masterpiece, best quality, 8k, ultra-detailed
        风格：如 realistic, anime, digital painting, cinematic lighting
        主体与数量：如 1girl, 1boy, group of people, animal, robot
        服饰与特征：如 blue dress, school uniform, ponytail, freckles, glasses
        动作与姿态：如 standing, sitting, looking at viewer, hand on hip
        背景与场景：如 indoors, sunset, cyberpunk city, forest, bedroom
        其它部分。
        生成的提示词仅为英文词组的集合，你无需给出任何中文描述，只需要返回提示词，不要有任何其他解释，无需换行
        严格围绕instruction，不要偏离主题
        你需要从参考图中提取画风，动作等，并将其融入到提示词中
        Return !!prompt 这里写你的图片提示词 !! ，不要有任何其它的解释
        请只返回以下格式：
        !!prompt 你的提示词 !!
        不要添加其它任何文字。
        你的提示词:
            
        '''
        #画风提取
    TEXT_PROMPT_STYLE = '''
        你是一个专业的AI绘画提示词工程师，用户的要求是：{instruction}
        这是用户的参考图，用英文提示词描述这个图片的画风
        提示词要用英文词组的形式给出，用逗号隔开，
        你需要遵循以下格式撰写提示词：
        风格：如 realistic, anime, digital painting, cinematic lighting
        你还需要提取出画风，并将其融入到提示词中，如油画，写实，厚涂，水彩，素描，等等
        生成的提示词仅为英文词组的集合，你无需给出任何中文描述，只需要返回提示词，不要有任何其他解释，无需换行
        严格围绕instruction，不要偏离主题
        你仅需要提取画风，无需提取其它任何部分
        Return !!prompt 这里写你的图片提示词 !! ，不要有任何其它的解释
        请只返回以下格式：
        !!prompt 你的提示词 !!
        不要添加其它任何文字。
        你的提示词:            
        '''

        #动作提取
    TEXT_PROMPT_ACTION = '''
        你是一个专业的AI绘画提示词工程师，用户的要求是：{instruction}
        这是用户的参考图，用英文提示词描述这个图片的动作
        提示词要用英文词组的形式给出，用逗号隔开，
        你需要遵循以下格式撰写提示词：
        动作与姿态：如 standing, sitting, looking at viewer, hand on hip
        你还需要提取出详细的动作，并将其融入到提示词中，如站立，坐下，看向镜头，手放在臀部，等等
        生成的提示词仅为英文词组的集合，你无需给出任何中文描述，只需要返回提示词，不要有任何其他解释，无需换行
        严格围绕instruction，不要偏离主题
        你仅需要提取动作，无需提取其它任何部分
        Return !!prompt 这里写你的图片提示词 !! ，不要有任何其它的解释
        请只返回以下格式：
        !!prompt 你的提示词 !!
        不要添加其它任何文字。
        你的提示词:
        '''

        #人物特征提取
    TEXT_PROMPT_TASK = '''
        你是一个专业的AI绘画提示词工程师，用户的要求是：{instruction}
        这是用户的参考图，用英文提示词描述这个图片的人物特征
        提示词要用英文词组的形式给出，用逗号隔开，
        你需要遵循以下格式撰写提示词：
        人物特征：如 长发，短发，长发飘飘，短发齐肩，眼睛颜色，头发颜色等等
        生成的提示词仅为英文词组的集合，你无需给出任何中文描述，只需要返回提示词，不要有任何其他解释，无需换行
        严格围绕instruction，不要偏离主题
        你仅需要提取人物特征，无需提取其它任何部分
        Return !!prompt 这里写你的图片提示词 !! ，不要有任何其它的解释
        请只返回以下格式：
        !!prompt 你的提示词 !!
        不要添加其它任何文字。
        你的提示词:
        '''
        #提示词合并
    Final_idea:str = '''
        生成图片{instruction}
        参考提示词：{ref_prompt}
        请将参考提示词融入到生成提示词中
        '''
    if temp_ref == "y":
        # 提示用户输入参考图和抽取范围
        ref_picture = input("Please type the url of the reference image")
        #检测url是否合法
        try:
            response = requests.get(ref_picture, timeout=5)
            if response.status_code != 200 or not response.headers.get("Content-Type", "").startswith("image"):
                print(f"图片 URL 无法访问或不是有效的图片，状态码：{response.status_code}")
                return
        except Exception as e:
            print(f"无法访问图片 URL，错误信息：{e}")
            return
        print("全部提取 : 1")
        print("画风提取 : 2")
        print("动作提取 : 3")
        print("人物特征提取 : 4")
        abstract_range = input("Press the number of abstract range (1-4): ")

        # 选择对应的 prompt 模板
        TEMPLATES = {
            "1": TEXT_PROMPT_ALL,
            "2": TEXT_PROMPT_STYLE,
            "3": TEXT_PROMPT_ACTION,
            "4": TEXT_PROMPT_TASK
        }
        template = TEMPLATES.get(abstract_range)
        
        if not template:
            print("无效选项，退出")
            return
        formatted_template = template.format(instruction=idea)
        # 提取并解析 prompt
        raw_prompt = image_to_text(ref_picture, formatted_template)
        ref_prompt = parse_prompt(raw_prompt)
        if not ref_prompt:
            print("参考提示词提取失败")
            return

        # 构造最终 idea
        idea = Final_idea.format(instruction=idea, ref_prompt=ref_prompt)

    else:
        # 不使用参考图
        print("Skipping reference...")
        # idea 保持原始输入
    logger.info(idea)
    team = Team()
    team.hire(
        [
            PromptWriter(),
            PictureGenerator(),
            SuggestionProvider(is_human=add_human)
        ]
    )
    team.invest(investment=investment)
    team.run_project(idea)
    await team.run(n_round=n_round)
if __name__ == "__main__":
    fire.Fire(main)